# -*- coding: utf-8 -*-
"""
🚀 物联网设备唤醒延迟测试平台（优化版）
功能：
  - 支持 60k TCP 长连接保活
  - 错峰心跳，避免风暴
  - 每分钟随机唤醒一批设备（UDP）
  - 记录唤醒延迟（UDP 发 → TCP 收）
  - 使用 uvloop 加速
  - 多进程分摊负载（可选）
  - 断线重连机制
  - 实时监控 & Excel 输出
"""

import asyncio
import socket
import struct
import time
import random
import string
import threading
import signal
from datetime import datetime
from typing import Dict, Optional, Set, List
import pandas as pd
import os
import logging
from concurrent.futures import ThreadPoolExecutor, as_completed

# ==================== 日志配置 ====================
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s | %(levelname)s | %(message)s',
    handlers=[
        logging.FileHandler(f"test_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# ==================== 尝试使用 uvloop ====================
try:
    import uvloop
    asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
    logger.info("🚀 使用 uvloop 加速异步事件循环")
except Exception as e:
    logger.warning(f"⚠️ uvloop 加载失败，使用原生 asyncio: {e}")

# ==================== 配置参数 ====================
DID_PREFIX = "IOTGDD"
TOTAL_DID_COUNT = 60000
BATCH_CONNECT_SIZE = 500      # 减小批次，更平滑
BATCH_CONNECT_INTERVAL = 0.5  # 批次间隔（秒）

TCP_HOST = "120.26.136.5"
TCP_PORT = 50000
TCP_HEARTBEAT_INTERVAL = 180  # 3分钟
TCP_CONNECT_TIMEOUT = 10
TCP_RECONNECT_DELAY = (5, 15)  # 重连间隔 5~15秒随机

UDP_HOST = "120.26.136.5"
UDP_PORT = 50001
UDP_BATCH_SIZE = 6000
UDP_TIMEOUT = 5
UDP_THREADS = 50

ENCRYPTION_KEY = "RSIOTGWakeupKey@"
ENCRYPTION_TYPE = 256

RUN_MINUTES = 10080  # 7天
OUTPUT_FILE = f"wakeup_latency_{datetime.now().strftime('%Y%m%d_%H%M%S')}.xlsx"

# ==================== 全局状态与线程安全锁 ====================
pending_wakeup: Dict[str, float] = {}
latency_results: Dict[str, Dict] = {}

# 🔐 线程安全锁
connect_lock = threading.Lock()
pending_lock = threading.Lock()
active_clients_lock = threading.Lock()
active_clients: Set['TCPDeviceClient'] = set()

# 统计
stats = {
    "connected": 0,
    "disconnected": 0,
    "reconnect_attempts": 0,
    "udp_sent": 0,
    "udp_success": 0,
    "responses": 0
}

# 事件：用于优雅关闭
shutdown_event = threading.Event()

# 导入加密
try:
    from string_enc import iPN_StringEnc
except ImportError:
    logger.error("❌ 请确保 string_enc.py 存在")
    exit(1)


# ==================== 工具函数 ====================
def generate_did_list() -> List[str]:
    chars = string.ascii_uppercase + string.digits
    return [
        f"IOTGDD-{i:06d}-{''.join(random.choices(chars, k=6))}"
        for i in range(TOTAL_DID_COUNT)
    ]


def create_udp_packet(did: str) -> bytes:
    payload = did.encode('utf-8')
    msg_len = len(payload)
    header = struct.pack(">III", 0xFFFFFFA1, 3, msg_len)
    return header + payload


def send_udp_wakeup(did: str) -> bool:
    packet = create_udp_packet(did)
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
            sock.settimeout(UDP_TIMEOUT)
            sock.sendto(packet, (UDP_HOST, UDP_PORT))
        return True
    except Exception:
        return False


# ==================== TCP 客户端（带重连）====================
class TCPDeviceClient:
    def __init__(self, did: str):
        self.did = did
        _, self.encrypted_data = iPN_StringEnc(ENCRYPTION_KEY, did, ENCRYPTION_TYPE)
        if isinstance(self.encrypted_data, str):
            self.encrypted_data = self.encrypted_data.encode('utf-8')
        self.reader: Optional[asyncio.StreamReader] = None
        self.writer: Optional[asyncio.StreamWriter] = None
        self.connected = False
        self.running = True

    async def connect_with_retry(self):
        while self.running and not shutdown_event.is_set():
            try:
                self.reader, self.writer = await asyncio.wait_for(
                    asyncio.open_connection(TCP_HOST, TCP_PORT),
                    timeout=TCP_CONNECT_TIMEOUT
                )
                self.connected = True

                with active_clients_lock:
                    active_clients.add(self)

                with connect_lock:
                    stats["connected"] += 1
                logger.info(f"✅ 连接成功 {self.did} | 当前活跃: {stats['connected']}")

                # 启动任务
                heartbeat_task = asyncio.create_task(self._start_heartbeat())
                listen_task = asyncio.create_task(self._listen())
                await asyncio.gather(heartbeat_task, listen_task)
                break  # 正常退出则不再重连

            except Exception as e:
                if not self.running:
                    break
                stats["disconnected"] += 1
                delay = random.uniform(*TCP_RECONNECT_DELAY)
                stats["reconnect_attempts"] += 1
                await asyncio.sleep(delay)  # 重连前等待

    async def _start_heartbeat(self):
        initial_delay = random.uniform(0, TCP_HEARTBEAT_INTERVAL)
        await asyncio.sleep(initial_delay)

        while self.connected and self.running and not shutdown_event.is_set():
            if self.writer and not self.writer.is_closing():
                try:
                    self.writer.write(self.encrypted_data)
                    await self.writer.drain()
                except Exception:
                    self.connected = False
                    break
            await asyncio.sleep(TCP_HEARTBEAT_INTERVAL)

    async def _listen(self):
        try:
            while self.connected and self.running and not shutdown_event.is_set():
                data = await self.reader.read(1024)
                if not data:
                    self.connected = False
                    break

                recv_time = time.time()
                with pending_lock:
                    if self.did in pending_wakeup:
                        send_time = pending_wakeup[self.did]
                        delay = recv_time - send_time
                        latency_results[self.did] = {
                            "DID": self.did,
                            "UDP_Send_Time": datetime.fromtimestamp(send_time).strftime("%Y-%m-%d %H:%M:%S.%f")[:-3],
                            "TCP_Recv_Time": datetime.fromtimestamp(recv_time).strftime("%Y-%m-%d %H:%M:%S.%f")[:-3],
                            "Latency_ms": round(delay * 1000, 2)
                        }
                        stats["responses"] += 1
                        logger.info(f"✅ 响应 {self.did} | 延迟: {delay * 1000:.2f}ms")
                        del pending_wakeup[self.did]
        except Exception:
            pass
        finally:
            self.connected = False
            with active_clients_lock:
                active_clients.discard(self)
            if self.writer and not self.writer.is_closing():
                self.writer.close()
                try:
                    await self.writer.wait_closed()
                except:
                    pass

    def stop(self):
        self.running = False
        self.connected = False


# ==================== UDP 唤醒任务 ====================
def trigger_udp_wakeup_batch():
    with active_clients_lock:
        live_dids = [client.did for client in active_clients if client.connected]

    if not live_dids:
        return

    selected_dids = random.choices(live_dids, k=min(UDP_BATCH_SIZE, len(live_dids)))
    current_time = time.time()

    with pending_lock:
        for did in selected_dids:
            pending_wakeup[did] = current_time

    send_time_str = datetime.fromtimestamp(current_time).strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
    logger.info(f"⏱️  发送 UDP 唤醒: {len(selected_dids)} 个")

    with ThreadPoolExecutor(max_workers=UDP_THREADS) as executor:
        futures = [executor.submit(send_udp_wakeup, did) for did in selected_dids]
        success = sum(1 for f in as_completed(futures) if f.result())

    stats["udp_sent"] += len(selected_dids)
    stats["udp_success"] += success
    logger.info(f"🚀 UDP 完成 | 发送: {len(selected_dids)} | 成功: {success}")


# ==================== 实时监控任务 ====================
async def monitor_status():
    """每 60 秒打印一次状态"""
    while not shutdown_event.is_set():
        await asyncio.sleep(60)
        with active_clients_lock:
            active_count = len([c for c in active_clients if c.connected])
        logger.info(
            f"📊 状态监控 | 活跃连接: {active_count} | "
            f"响应: {stats['responses']} | "
            f"UDP成功: {stats['udp_success']} | "
            f"重连: {stats['reconnect_attempts']}"
        )


# ==================== 主程序 ====================
async def main():
    logger.info(f"🚀 物联网设备唤醒延迟测试平台启动")
    logger.info(f"  TCP 连接: {TOTAL_DID_COUNT} | 分批: {BATCH_CONNECT_SIZE}")
    logger.info(f"  UDP 唤醒: {UDP_BATCH_SIZE}/分钟 | 运行: {RUN_MINUTES}分钟")

    did_list = generate_did_list()
    tcp_clients = [TCPDeviceClient(did) for did in did_list]

    # 启动监控
    asyncio.create_task(monitor_status())

    # 分批连接
    logger.info("🟡 正在分批建立 TCP 连接...")
    for i in range(0, len(tcp_clients), BATCH_CONNECT_SIZE):
        batch = tcp_clients[i:i + BATCH_CONNECT_SIZE]
        for client in batch:
            asyncio.create_task(client.connect_with_retry())
        await asyncio.sleep(BATCH_CONNECT_INTERVAL)

    logger.info(f"✅ 已提交 {len(tcp_clients)} 个连接任务")

    # 主唤醒循环
    start_time = time.time()
    try:
        for minute in range(RUN_MINUTES):
            if shutdown_event.is_set():
                break
            sleep_time = 60 - (time.time() - start_time) % 60
            await asyncio.sleep(max(sleep_time, 0.1))
            trigger_udp_wakeup_batch()
    except KeyboardInterrupt:
        logger.info("🛑 用户中断")
    finally:
        shutdown_event.set()
        logger.info("🛑 正在关闭所有连接...")
        for client in tcp_clients:
            client.stop()
        save_results()


def save_results():
    if not latency_results:
        logger.warning("⚠️ 未记录到任何延迟数据")
        return

    df = pd.DataFrame(list(latency_results.values()))
    df = df.sort_values("Latency_ms")

    avg = df["Latency_ms"].mean()
    p95 = df["Latency_ms"].quantile(0.95)
    p99 = df["Latency_ms"].quantile(0.99)

    with pd.ExcelWriter(OUTPUT_FILE, engine='openpyxl') as writer:
        df.to_excel(writer, sheet_name='Latency', index=False)
        stats_df = pd.DataFrame({
            "Metric": ["Average (ms)", "P95 (ms)", "P99 (ms)", "Total Samples"],
            "Value": [round(avg, 2), round(p95, 2), round(p99, 2), len(df)]
        })
        stats_df.to_excel(writer, sheet_name='Summary', index=False)

    logger.info(f"📊 结果已保存: {OUTPUT_FILE}")
    logger.info(f"📈 平均延迟: {avg:.2f}ms | P95: {p95:.2f}ms | P99: {p99:.2f}ms")


# ==================== 优雅关闭 ====================
def signal_handler(signum, frame):
    logger.info(f"🛑 收到信号 {signum}，正在优雅关闭...")
    shutdown_event.set()


if __name__ == "__main__":
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    asyncio.run(main())